package com.jlpay.open.jlpay.sdk.java.utils.gm;

import java.math.BigInteger;
import java.security.SecureRandom;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECKeyParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.math.ec.ECConstants;
import org.bouncycastle.math.ec.ECFieldElement;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.BigIntegers;

public class Sm2EngineExtend {
    private final Digest digest;
    private boolean forEncryption;
    private ECKeyParameters ecKey;
    private ECDomainParameters ecParams;
    private int curveLength;
    private SecureRandom random;
    private int cipherMode;
    public static final int CIPHERMODE_BC = 0;
    public static final int CIPHERMODE_NORM = 1;

    public Sm2EngineExtend() {
        this(new SM3Digest());
    }

    public Sm2EngineExtend(Digest digest) {
        this.digest = digest;
    }

    public void setCipherMode(int cipherMode) {
        this.cipherMode = cipherMode;
    }

    public void init(boolean forEncryption, CipherParameters param) {
        this.init(forEncryption, 1, param);
    }

    public void init(boolean forEncryption, int cipherMode, CipherParameters param) {
        this.forEncryption = forEncryption;
        this.cipherMode = cipherMode;
        if (forEncryption) {
            ParametersWithRandom rParam = (ParametersWithRandom)param;
            this.ecKey = (ECKeyParameters)rParam.getParameters();
            this.ecParams = this.ecKey.getParameters();
            ECPoint s = ((ECPublicKeyParameters)this.ecKey).getQ().multiply(this.ecParams.getH());
            if (s.isInfinity()) {
                throw new IllegalArgumentException("invalid key: [h]Q at infinity");
            }

            this.random = rParam.getRandom();
        } else {
            this.ecKey = (ECKeyParameters)param;
            this.ecParams = this.ecKey.getParameters();
        }

        this.curveLength = (this.ecParams.getCurve().getFieldSize() + 7) / 8;
    }

    public byte[] processBlock(byte[] in, int inOff, int inLen) throws InvalidCipherTextException {
        return this.forEncryption ? this.encrypt(in, inOff, inLen) : this.decrypt(in, inOff, inLen);
    }

    private byte[] encrypt(byte[] in, int inOff, int inLen) {
        byte[] c2 = new byte[inLen];
        System.arraycopy(in, inOff, c2, 0, c2.length);

        byte[] c1;
        ECPoint kpb;
        do {
            BigInteger k = this.nextk();
            ECPoint c1p = this.ecParams.getG().multiply(k).normalize();
            c1 = c1p.getEncoded(false);
            kpb = ((ECPublicKeyParameters)this.ecKey).getQ().multiply(k).normalize();
            this.kdf(this.digest, kpb, c2);
        } while(this.notEncrypted(c2, in, inOff));

        this.addFieldElement(this.digest, kpb.getAffineXCoord());
        this.digest.update(in, inOff, inLen);
        this.addFieldElement(this.digest, kpb.getAffineYCoord());
        byte[] c3 = new byte[this.digest.getDigestSize()];
        this.digest.doFinal(c3, 0);
        return this.cipherMode == 1 ? Arrays.concatenate(c1, c3, c2) : Arrays.concatenate(c1, c2, c3);
    }

    private byte[] decrypt(byte[] in, int inOff, int inLen) throws InvalidCipherTextException {
        byte[] c1 = new byte[this.curveLength * 2 + 1];
        System.arraycopy(in, inOff, c1, 0, c1.length);
        ECPoint c1p = this.ecParams.getCurve().decodePoint(c1);
        ECPoint s = c1p.multiply(this.ecParams.getH());
        if (s.isInfinity()) {
            throw new InvalidCipherTextException("[h]C1 at infinity");
        } else {
            c1p = c1p.multiply(((ECPrivateKeyParameters)this.ecKey).getD()).normalize();
            byte[] c2 = new byte[inLen - c1.length - this.digest.getDigestSize()];
            if (this.cipherMode == 0) {
                System.arraycopy(in, inOff + c1.length, c2, 0, c2.length);
            } else {
                System.arraycopy(in, inOff + c1.length + this.digest.getDigestSize(), c2, 0, c2.length);
            }

            this.kdf(this.digest, c1p, c2);
            this.addFieldElement(this.digest, c1p.getAffineXCoord());
            this.digest.update(c2, 0, c2.length);
            this.addFieldElement(this.digest, c1p.getAffineYCoord());
            byte[] c3 = new byte[this.digest.getDigestSize()];
            this.digest.doFinal(c3, 0);
            int check = 0;
            int i;
            if (this.cipherMode == 0) {
                for(i = 0; i != c3.length; ++i) {
                    check |= c3[i] ^ in[c1.length + c2.length + i];
                }
            } else {
                for(i = 0; i != c3.length; ++i) {
                    check |= c3[i] ^ in[c1.length + i];
                }
            }

            this.clearBlock(c1);
            this.clearBlock(c3);
            if (check != 0) {
                this.clearBlock(c2);
                throw new InvalidCipherTextException("invalid cipher text");
            } else {
                return c2;
            }
        }
    }

    private boolean notEncrypted(byte[] encData, byte[] in, int inOff) {
        for(int i = 0; i != encData.length; ++i) {
            if (encData[i] != in[inOff]) {
                return false;
            }
        }

        return true;
    }

    private void kdf(Digest digest, ECPoint c1, byte[] encData) {
        int ct = 1;
        int v = digest.getDigestSize();
        byte[] buf = new byte[digest.getDigestSize()];
        int off = 0;

        for(int i = 1; i <= (encData.length + v - 1) / v; ++i) {
            this.addFieldElement(digest, c1.getAffineXCoord());
            this.addFieldElement(digest, c1.getAffineYCoord());
            digest.update((byte)(ct >> 24));
            digest.update((byte)(ct >> 16));
            digest.update((byte)(ct >> 8));
            digest.update((byte)ct);
            digest.doFinal(buf, 0);
            if (off + buf.length < encData.length) {
                this.xor(encData, buf, off, buf.length);
            } else {
                this.xor(encData, buf, off, encData.length - off);
            }

            off += buf.length;
            ++ct;
        }

    }

    private void xor(byte[] data, byte[] kdfOut, int dOff, int dRemaining) {
        for(int i = 0; i != dRemaining; ++i) {
            data[dOff + i] ^= kdfOut[i];
        }

    }

    private BigInteger nextk() {
        int qBitLength = this.ecParams.getN().bitLength();

        BigInteger k;
        do {
            do {
                k = new BigInteger(qBitLength, this.random);
            } while(k.equals(ECConstants.ZERO));
        } while(k.compareTo(this.ecParams.getN()) >= 0);

        return k;
    }

    private void addFieldElement(Digest digest, ECFieldElement v) {
        byte[] p = BigIntegers.asUnsignedByteArray(this.curveLength, v.toBigInteger());
        digest.update(p, 0, p.length);
    }

    private void clearBlock(byte[] block) {
        java.util.Arrays.fill(block, (byte)0);
    }
}
